import logging
import time
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass
from threading import RLock
from contextlib import contextmanager

logger = logging.getLogger(__name__)


@dataclass
class TransactionContext:
    """事务上下文类"""
    transaction_id: str
    start_time: float
    operations: List[Dict[str, Any]]
    status: str = "active"
    lock: RLock = None

    def __post_init__(self):
        self.lock = self.lock or RLock()


class DataConsistencyManager:
    """
    数据一致性管理器
    负责管理分布式系统中的数据一致性，包括两阶段提交、数据同步和错误恢复
    """
    
    def __init__(self):
        self.active_transactions: Dict[str, TransactionContext] = {}
        self.transaction_lock = RLock()
        self.max_transaction_timeout = 300  # 5分钟
        self.sync_retry_count = 3
        self.sync_retry_delay = 2
    
    def begin_transaction(self, transaction_id: str) -> TransactionContext:
        """
        开始一个新事务
        """
        with self.transaction_lock:
            if transaction_id in self.active_transactions:
                raise ValueError(f"事务ID已存在: {transaction_id}")
            
            tx = TransactionContext(
                transaction_id=transaction_id,
                start_time=time.time(),
                operations=[]
            )
            self.active_transactions[transaction_id] = tx
            logger.info(f"事务开始: {transaction_id}")
            return tx
    
    def add_operation(self, transaction_id: str, operation: Dict[str, Any]) -> None:
        """
        向事务中添加操作
        """
        with self.transaction_lock:
            tx = self.active_transactions.get(transaction_id)
            if not tx:
                raise ValueError(f"事务不存在: {transaction_id}")
            
        with tx.lock:
            if tx.status != "active":
                raise ValueError(f"事务已关闭: {transaction_id}")
            
            # 检查事务是否超时
            if time.time() - tx.start_time > self.max_transaction_timeout:
                tx.status = "timeout"
                raise TimeoutError(f"事务超时: {transaction_id}")
            
            tx.operations.append(operation)
            logger.debug(f"事务 {transaction_id} 添加操作: {operation}")
    
    def commit_transaction(self, transaction_id: str) -> bool:
        """
        提交事务（两阶段提交）
        1. 准备阶段：验证所有操作是否可以执行
        2. 提交阶段：实际执行所有操作
        """
        try:
            with self.transaction_lock:
                tx = self.active_transactions.get(transaction_id)
                if not tx:
                    raise ValueError(f"事务不存在: {transaction_id}")
            
            with tx.lock:
                if tx.status != "active":
                    raise ValueError(f"事务状态错误: {transaction_id}, 状态: {tx.status}")
                
                tx.status = "preparing"
                logger.info(f"事务准备提交: {transaction_id}")
                
                # 1. 准备阶段 - 验证所有操作
                can_commit = self._prepare_transaction(tx)
                
                if not can_commit:
                    tx.status = "failed_prepare"
                    logger.error(f"事务准备失败: {transaction_id}")
                    return False
                
                # 2. 提交阶段 - 执行所有操作
                tx.status = "committing"
                success = self._commit_operations(tx)
                
                if success:
                    tx.status = "committed"
                    logger.info(f"事务提交成功: {transaction_id}")
                    
                    # 清理已完成的事务
                    with self.transaction_lock:
                        del self.active_transactions[transaction_id]
                    
                    return True
                else:
                    tx.status = "failed_commit"
                    logger.error(f"事务提交失败: {transaction_id}")
                    return False
                    
        except Exception as e:
            logger.exception(f"事务提交异常: {transaction_id}")
            return False
    
    def rollback_transaction(self, transaction_id: str) -> bool:
        """
        回滚事务
        """
        try:
            with self.transaction_lock:
                tx = self.active_transactions.get(transaction_id)
                if not tx:
                    return True  # 事务不存在，视为已回滚
            
            with tx.lock:
                if tx.status in ["committed", "rolled_back"]:
                    return True  # 事务已结束，无需回滚
                
                tx.status = "rolling_back"
                logger.info(f"事务开始回滚: {transaction_id}")
                
                # 尝试回滚已执行的操作
                self._rollback_operations(tx)
                
                tx.status = "rolled_back"
                logger.info(f"事务回滚完成: {transaction_id}")
                
                # 清理事务
                with self.transaction_lock:
                    del self.active_transactions[transaction_id]
                
                return True
                
        except Exception as e:
            logger.exception(f"事务回滚异常: {transaction_id}")
            return False
    
    def _prepare_transaction(self, tx: TransactionContext) -> bool:
        """
        准备事务 - 验证所有操作是否可以执行
        """
        try:
            for operation in tx.operations:
                # 这里应该根据不同的操作类型进行验证
                # 例如：检查记录是否存在，检查权限等
                pass
            return True
        except Exception as e:
            logger.error(f"事务准备验证失败: {tx.transaction_id}, 错误: {str(e)}")
            return False
    
    def _commit_operations(self, tx: TransactionContext) -> bool:
        """
        提交操作 - 实际执行所有操作
        """
        executed_operations = []
        try:
            for i, operation in enumerate(tx.operations):
                # 根据操作类型执行实际操作
                success = self._execute_operation(operation)
                if success:
                    executed_operations.append((i, operation))
                else:
                    logger.error(f"操作执行失败: {operation}")
                    return False
            return True
        except Exception as e:
            logger.error(f"操作执行异常: {str(e)}")
            return False
    
    def _rollback_operations(self, tx: TransactionContext) -> bool:
        """
        回滚操作 - 撤销已执行的操作
        """
        try:
            # 逆序回滚已执行的操作
            for operation in reversed(tx.operations):
                # 根据操作类型执行对应的回滚操作
                self._rollback_operation(operation)
            return True
        except Exception as e:
            logger.error(f"操作回滚异常: {str(e)}")
            return False
    
    def _execute_operation(self, operation: Dict[str, Any]) -> bool:
        """
        执行单个操作
        """
        try:
            # 这里应该根据不同的操作类型执行实际操作
            # 例如：insert, update, delete 等
            operation_type = operation.get("type")
            logger.debug(f"执行操作: {operation_type}")
            return True
        except Exception as e:
            logger.error(f"操作执行失败: {operation}, 错误: {str(e)}")
            return False
    
    def _rollback_operation(self, operation: Dict[str, Any]) -> bool:
        """
        回滚单个操作
        """
        try:
            # 这里应该根据不同的操作类型执行对应的回滚操作
            operation_type = operation.get("type")
            logger.debug(f"回滚操作: {operation_type}")
            return True
        except Exception as e:
            logger.error(f"操作回滚失败: {operation}, 错误: {str(e)}")
            return False
    
    def sync_data(self, source: str, target: str, data: List[Dict[str, Any]], 
                 sync_type: str = "full") -> Dict[str, Any]:
        """
        同步数据
        :param source: 数据源标识
        :param target: 目标数据源标识
        :param data: 要同步的数据
        :param sync_type: 同步类型: "full" 或 "incremental"
        :return: 同步结果
        """
        sync_id = f"sync_{source}_to_{target}_{int(time.time())}"
        logger.info(f"开始数据同步: {sync_id}, 数据量: {len(data)}, 类型: {sync_type}")
        
        retry_count = 0
        while retry_count <= self.sync_retry_count:
            try:
                if retry_count > 0:
                    logger.info(f"重试数据同步: {sync_id}, 重试次数: {retry_count}")
                    time.sleep(self.sync_retry_delay * retry_count)
                
                # 执行同步操作
                result = self._perform_sync(source, target, data, sync_type)
                
                if result["success"]:
                    logger.info(f"数据同步成功: {sync_id}, 同步记录数: {result.get('synced_count', 0)}")
                    return result
                
            except Exception as e:
                logger.error(f"数据同步异常: {sync_id}, 错误: {str(e)}")
            
            retry_count += 1
        
        # 同步失败
        logger.error(f"数据同步失败: {sync_id}, 已达到最大重试次数")
        return {
            "success": False,
            "error": "数据同步失败，已达到最大重试次数",
            "sync_id": sync_id,
            "retry_count": retry_count
        }
    
    def _perform_sync(self, source: str, target: str, data: List[Dict[str, Any]], 
                     sync_type: str) -> Dict[str, Any]:
        """
        执行实际的同步操作
        """
        # 这里应该实现实际的数据同步逻辑
        # 例如：在PostgreSQL和Milvus之间同步数据
        synced_count = 0
        failed_items = []
        
        for item in data:
            try:
                # 执行单个记录的同步
                # 这里应该根据source和target的类型实现不同的同步逻辑
                synced_count += 1
            except Exception as e:
                failed_items.append({
                    "item": item,
                    "error": str(e)
                })
        
        return {
            "success": len(failed_items) == 0,
            "synced_count": synced_count,
            "failed_count": len(failed_items),
            "failed_items": failed_items
        }
    
    def recover_from_failure(self, failure_info: Dict[str, Any]) -> Dict[str, Any]:
        """
        从故障中恢复
        :param failure_info: 故障信息
        :return: 恢复结果
        """
        recovery_id = f"recovery_{int(time.time())}"
        logger.info(f"开始故障恢复: {recovery_id}, 故障类型: {failure_info.get('failure_type')}")
        
        try:
            failure_type = failure_info.get('failure_type')
            
            if failure_type == "transaction_failure":
                # 事务失败恢复
                transaction_id = failure_info.get('transaction_id')
                return self._recover_transaction(transaction_id)
            
            elif failure_type == "sync_failure":
                # 同步失败恢复
                source = failure_info.get('source')
                target = failure_info.get('target')
                failed_items = failure_info.get('failed_items', [])
                return self._recover_sync(source, target, failed_items)
            
            elif failure_type == "system_crash":
                # 系统崩溃恢复
                return self._recover_system_state()
            
            else:
                logger.error(f"未知故障类型: {failure_type}")
                return {"success": False, "error": "未知故障类型"}
                
        except Exception as e:
            logger.exception(f"故障恢复异常: {recovery_id}")
            return {"success": False, "error": str(e)}
    
    def _recover_transaction(self, transaction_id: str) -> Dict[str, Any]:
        """
        恢复失败的事务
        """
        try:
            with self.transaction_lock:
                tx = self.active_transactions.get(transaction_id)
                if not tx:
                    logger.info(f"事务不存在，无需恢复: {transaction_id}")
                    return {"success": True, "message": "事务不存在或已完成"}
            
            # 尝试回滚事务
            success = self.rollback_transaction(transaction_id)
            return {
                "success": success,
                "transaction_id": transaction_id
            }
            
        except Exception as e:
            logger.error(f"事务恢复失败: {transaction_id}, 错误: {str(e)}")
            return {"success": False, "error": str(e)}
    
    def _recover_sync(self, source: str, target: str, failed_items: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        恢复失败的同步操作
        """
        data_to_retry = [item["item"] for item in failed_items]
        result = self.sync_data(source, target, data_to_retry, "incremental")
        return result
    
    def _recover_system_state(self) -> Dict[str, Any]:
        """
        系统崩溃后恢复状态
        """
        try:
            # 处理所有活跃事务
            with self.transaction_lock:
                active_txs = list(self.active_transactions.values())
            
            recovery_results = []
            for tx in active_txs:
                result = self._recover_transaction(tx.transaction_id)
                recovery_results.append({
                    "transaction_id": tx.transaction_id,
                    "result": result
                })
            
            # 执行数据一致性检查
            consistency_check = self._check_data_consistency()
            
            return {
                "success": all(r["result"]["success"] for r in recovery_results),
                "recovered_transactions": len(recovery_results),
                "consistency_check": consistency_check
            }
            
        except Exception as e:
            logger.error(f"系统状态恢复失败: {str(e)}")
            return {"success": False, "error": str(e)}
    
    def _check_data_consistency(self) -> Dict[str, Any]:
        """
        检查数据一致性
        """
        # 这里应该实现数据一致性检查逻辑
        # 例如：比较不同数据源之间的数据是否一致
        logger.info("执行数据一致性检查")
        return {
            "consistent": True,
            "message": "数据一致性检查通过"
        }


# 创建全局数据一致性管理器实例
data_consistency_manager = DataConsistencyManager()


@contextmanager
def transaction(transaction_id: str):
    """
    事务上下文管理器
    """
    manager = data_consistency_manager
    
    try:
        # 开始事务
        manager.begin_transaction(transaction_id)
        
        # 提供事务管理器给上下文
        yield manager
        
        # 提交事务
        success = manager.commit_transaction(transaction_id)
        if not success:
            raise Exception(f"事务提交失败: {transaction_id}")
            
    except Exception as e:
        # 发生异常，回滚事务
        try:
            manager.rollback_transaction(transaction_id)
        except:
            pass
        raise e
